'''
后面几个提取器接收的都是packet列表
直接减少多次读取同一文件的开销

本文件的功能是提取同一个流中pkt的特征
流的包数量
包字节 总数 最大值 最小值 均值 标准差
包间间隔时间 最大值 最小值 均值 标准差
流持续时间
熵
'''
'''
外部调用接口
生成实例对象，传入packet列表，调用Extract_Features方法，即可得到流的特征
值可以通过print打印出来，也可以直接获取写入
'''

import numpy as np
import math
from scapy.all import *

class FlowFeatureExtractor:
    def __init__(self, packet_list):
        self.pkts = packet_list
        
        # 待保存的值
        self.pkt_num = 0
        
        self.pkt_total_bytes = 0
        self.pkt_max_bytes = 0
        self.pkt_min_bytes = 0
        self.pkt_mean_bytes = 0
        self.pkt_std_bytes = 0
        
        self.pkt_interval_time = 0
        self.pkt_max_interval_time = 0
        self.pkt_min_interval_time = 0
        self.pkt_mean_interval_time = 0
        self.pkt_std_interval_time = 0
        
        self.flow_duration = 0
        
        self.pkt_entropy = 0
        
        # 计算特征
        self.Extract_Features()
    
    def Extract_Features(self, filed = "TCP"):
        # 计算包数量
        self.pkt_num = len(self.pkts)
        
        # 计算包字节总数、最大值、最小值、均值、标准差
        try:
            self.pkt_total_bytes = sum([len(pkt[filed].payload) for pkt in self.pkts])
            self.pkt_max_bytes = max([len(pkt[filed].payload) for pkt in self.pkts])
            self.pkt_min_bytes = min([len(pkt[filed].payload) for pkt in self.pkts])
            self.pkt_mean_bytes = self.pkt_total_bytes / self.pkt_num
            self.pkt_std_bytes = np.std([len(pkt[filed].payload) for pkt in self.pkts])
        except:
            pass
        
        # 计算包间间隔时间、最大值、最小值、均值、标准差
        try:
            interval_times = [self.pkts[i].time - self.pkts[i-1].time for i in range(1, len(self.pkts))]
            self.pkt_interval_time = sum(interval_times)
            self.pkt_max_interval_time = max(interval_times)
            self.pkt_min_interval_time = min(interval_times)
            self.pkt_mean_interval_time = sum(interval_times) / self.pkt_num
        except:
            pass
        
        # 原本值是numpy.int64的列表，需要转换为np.array之后，再强行转为float才能使用np.std计算标准差
        self.pkt_std_interval_time = np.std(np.array(interval_times).astype(float))
        
        # 计算流持续时间
        self.flow_duration = self.pkts[-1].time - self.pkts[0].time
        
        # 计算熵
        self.pkt_entropy = self.Calculate_Entropy(self.pkts)
    
    def Calculate_Entropy(self, packet_list):
        # 熵计算函数
        # 输入：packet_list，一个pkt列表
        # 输出：熵值
        # 针对每个数据包，计算每个字节出现的概率，然后计算熵值
        # 然后返回整个流的平均熵值
        entropy = 0
        for pkt in packet_list:
            byte_list = [0] * 256
            for byte in raw(pkt.payload):
                byte_list[byte] += 1
            
            N = sum(byte_list)
            for ni in byte_list:
                if ni > 0:
                    entropy -= ni / N * np.log2(ni / N)
        
        entropy /= len(packet_list)
        return entropy
    
    def __str__(self) -> str:
        return f"Packet Number: {self.pkt_num}\n" \
               f"Packet Total Bytes: {self.pkt_total_bytes}\n" \
               f"Packet Max Bytes: {self.pkt_max_bytes}\n" \
               f"Packet Min Bytes: {self.pkt_min_bytes}\n" \
               f"Packet Mean Bytes: {self.pkt_mean_bytes}\n" \
               f"Packet Std Bytes: {self.pkt_std_bytes}\n" \
               f"Packet Interval Time: {self.pkt_interval_time}\n" \
               f"Packet Max Interval Time: {self.pkt_max_interval_time}\n" \
               f"Packet Min Interval Time: {self.pkt_min_interval_time}\n" \
               f"Packet Mean Interval Time: {self.pkt_mean_interval_time}\n" \
               f"Packet Std Interval Time: {self.pkt_std_interval_time}\n" \
               f"Flow Duration: {self.flow_duration}\n" \
               f"Packet Entropy: {self.pkt_entropy}\n"
               
    def Return_Features(self):
        # 返回特征值列表
        return [self.pkt_num, self.pkt_total_bytes, self.pkt_max_bytes, self.pkt_min_bytes, self.pkt_mean_bytes, self.pkt_std_bytes,
                self.pkt_interval_time, self.pkt_max_interval_time, self.pkt_min_interval_time, self.pkt_mean_interval_time, self.pkt_std_interval_time,
                self.flow_duration, self.pkt_entropy]
    
if __name__ == '__main__':
    pcap_file = "./processed_dataset/browsing_flow_3.pcap"
    packets = sniff(offline=pcap_file, filter="")
    Test_obj = FlowFeatureExtractor(packet_list=packets)
    print(Test_obj)
    